# -*- coding: utf-8 -*-
from sqlalchemy import Column, Integer, BigInteger, String, JSON, DateTime
from sqlalchemy.orm import Session
from sqlalchemy.orm.query import Query

from datetime import datetime


from sqlalchemy.ext.declarative import declarative_base
from core.database import sqlite, mysql

DB_TYPE = "mysql"

async def register_db(dsn: str, echo: bool=False):
    global DB_TYPE
    if dsn.startswith("mysql"):
        DB_TYPE = "mysql"
        await mysql.register_db(dsn, echo)
    elif dsn.startswith("sqlite"):
        DB_TYPE = "sqlite"
        await sqlite.register_db(dsn, echo)
    else:
        raise Exception("Unknown database type")


def get_db():
    if DB_TYPE == "mysql":
        session = mysql.SessionLocal()
    elif DB_TYPE == "sqlite":
        session = sqlite.SessionLocal()
    else:
        raise Exception("Unknown database type")
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()
        
Base = declarative_base()

async def init_db():
    if DB_TYPE == "mysql":
        engine = mysql.engine
    elif DB_TYPE == "sqlite":
        engine = sqlite.engine
    else:
        raise Exception("Unknown database type")
    Base.metadata.create_all(bind=engine)

# from core.database.sqlite import register_db
# from core.database.sqlite import Base
# from core.database.sqlite import get_db

# from core.database.mysql import register_db
# from core.database.mysql import Base
# from core.database.mysql import get_db

# 根据表名获取模型类的函数
def get_model_class_by_table_name(table_name: str):
    for model_class in Base.registry._class_registry.values():
        if hasattr(model_class, '__tablename__') and model_class.__tablename__ == table_name:
            return model_class
    return None  # 如果没有找到匹配的模型类，返回 None

class TimestampMixin():
    extra = Column(JSON, default=dict)
    remark =  Column(String(256))
    create_source =  Column(String(256))
    create_time = Column(DateTime, default=datetime.now)
    create_user =  Column(BigInteger, default=-1)
    update_time = Column(DateTime, default=datetime.now)
    update_user =  Column(BigInteger, default=-1)
    del_flag =  Column(Integer, default=0)
    
    def __repr__(self):
        return "[%s%s]" % (self.__tablename__, str({c.name: getattr(self, c.name) for c in self.__table__.columns}))

    def single_to_dict(self):
        return {c.name: getattr(self, c.name).strftime("%Y-%m-%d %H:%M:%S") if isinstance(getattr(self, c.name), datetime) else getattr(self, c.name) for c in self.__table__.columns}

    @classmethod
    def filter(self, session: Session, filter: dict={}, like: dict={}, in_: dict={})-> Query:
        # table_ = self.__class__
        table_ = self
        if table_ is None:
            return [], 0
        filter_ = [getattr(table_, "del_flag") == 0]
        
        filter_ += [getattr(table_, key) == val for key, val in filter.items() if val is not None] if filter else []
        filter_ += [getattr(table_, key).like(f'%{val}%') for key, val in like.items() if val is not None] if like else []
        filter_ += [getattr(table_, key).in_(val) for key, val in in_.items() if isinstance(val, list)] if in_ else []
        return session.query(table_).filter(*filter_)

    @classmethod
    def order_by(self, order_by: dict={}):
        table_ = self
        return [table_.__dict__[key] if rev else table_.__dict__[key].desc() for key, rev in order_by.items()]

    def add(self, session: Session, commit: bool=False, flush: bool = True):
        session.add(self)
        if commit:
            session.commit()
        elif flush:
            session.flush()
        return self
        
    @classmethod
    def add_all(self, session: Session, data: list, commit: bool=False, flush: bool = True):
        session.add_all(data)
        if flush:
            session.flush()
        if commit:
            session.commit()
    
    def delete(self, session: Session, commit: bool=False):
        self.del_flag = 1
        if commit:
            session.commit()
        elif flush:
            session.flush()
        return self
    
    def update(self, session: Session, data: dict={}, commit: bool=False, flush: bool = True):
        for key, val in data.items():
            if val is None or key not in self.__table__.columns: continue
            setattr(self, key, val)
        if commit:
            session.commit()
        elif flush:
            session.flush()
        return self

